{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 7: Train NicheTrans on 10x Xenium data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, time, datetime, warnings\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.optim import lr_scheduler\n",
    "\n",
    "from model.nicheTrans_img import *\n",
    "from datasets.data_manager_breast_cancer import Breast_cancer\n",
    "\n",
    "from utils.utils import *\n",
    "from utils.utils_training_breast_cancer import train, test\n",
    "from utils.utils_dataloader import *\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize the args and fix seeds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==========\n",
      "Args:Namespace(adata_path='/home/wzk/ST_data/2024_NicheTrans_upload/2023_nc_10x_breast_cancer/HBC_rep1_cell_nucleus_3channel_strength_mean.h5ad', coordinate_path='/home/wzk/ST_data/2023_nc_Xenium_breast/In_situ_sample_1_replicate_1/outs/cells.csv.gz', ct_path='/home/wzk/ST_data/2023_nc_Xenium_breast/Cell_Barcode_Type_Matrices.xlsx', dropout_rate=0.2, eval_step=1, gamma=0.1, gpu_devices='0', lr=0.0003, max_epoch=40, noise_rate=0.2, optimizer='adam', seed=1, stepsize=20, test_batch=32, train_batch=32, weight_decay=0.0005, workers=4)\n",
      "==========\n"
     ]
    }
   ],
   "source": [
    "%run ./args/args_breast_cancer.py\n",
    "args = args\n",
    "\n",
    "set_seed(args.seed)\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices\n",
    "\n",
    "print(\"==========\\nArgs:{}\\n==========\".format(args))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize dataloaders and NicheTrans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------Calculating spatial graph...\n",
      "The graph contains 1185564 edges, 98797 cells.\n",
      "12.0000 neighbors per cell on average.\n",
      "------Calculating spatial graph...\n",
      "The graph contains 827796 edges, 68983 cells.\n",
      "12.0000 neighbors per cell on average.\n",
      "=> AD Mouse loaded\n",
      "Dataset statistics:\n",
      "  ------------------------------\n",
      "  subset   | # num | \n",
      "  ------------------------------\n",
      "  train    |  98797 spots, 98659 positive CD20, 84043 positive HER2 \n",
      "  test     |  68983 spots, 67600 positive CD20, 36904 positive HER2 \n",
      "  ------------------------------\n"
     ]
    }
   ],
   "source": [
    "# create the dataloaders\n",
    "dataset = Breast_cancer(adata_path=args.adata_path, coordinate_path=args.coordinate_path, ct_path=args.ct_path)\n",
    "trainloader, testloader = breast_cancer_dataloader(args, dataset)\n",
    "\n",
    "# create the model\n",
    "source_dimension, target_dimension = dataset.rna_length, dataset.protein_length\n",
    "model = NicheTrans(source_length=source_dimension, target_length=target_dimension, noise_rate=args.noise_rate, dropout_rate=args.dropout_rate)\n",
    "model = nn.DataParallel(model).cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize loss function (criterion) and optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.MSELoss()\n",
    "\n",
    "if args.optimizer == 'adam':\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)\n",
    "elif args.optimizer == 'SGD':\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=args.lr)\n",
    "else:\n",
    "    print('unexpected optimizer')\n",
    "\n",
    "if args.stepsize > 0:\n",
    "    scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model training and testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "\n",
    "for epoch in range(args.max_epoch):\n",
    "    last_epoch = epoch + 1 == args.max_epoch\n",
    "\n",
    "    print(\"==> Epoch {}/{}\".format(epoch+1, args.max_epoch))\n",
    "    \n",
    "    ################\n",
    "    train(args, model, criterion, optimizer, trainloader, dataset.target_panel)\n",
    "    if args.stepsize > 0: scheduler.step()\n",
    "    \n",
    "    if (epoch+1) % args.eval_step == 0:\n",
    "        pearson = test(args, model, testloader, dataset.target_panel, last_epoch)\n",
    "\n",
    "    if last_epoch==True:\n",
    "        torch.save(model.state_dict(), 'NicheTrans_breast_cancer_last.pth')\n",
    "    ################\n",
    "\n",
    "elapsed = round(time.time() - start_time)\n",
    "elapsed = str(datetime.timedelta(seconds=elapsed))\n",
    "print(\"Finished. Total elapsed time (h:m:s): {}\".format(elapsed))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch_zk",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}